import os, sys
import os.path as osp
import base64, copy
import io
import time
from PIL import Image
import numpy as np
import cv2
import json
import requests

from rec.handwrite_recognition  import handwrite_rec_func, pagenum_rec_func
from rec.seal_recognition  import seal_recognition_func
from rec.seal_cls import seal_class_func
from tool import filesystem

class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, bytes):
            return obj.decode("utf-8")
        else:
            return super(MyEncoder, self).default(obj)
        


def iou_3(rect1, rect2):
    y0 = np.max([rect1[1],rect2[1]])
    y1 = np.min([rect1[1] + rect1[3],rect2[1]+rect2[3]])

    x0 = np.max([rect1[0],rect2[0]])
    x1 = np.min([rect1[0] + rect1[2],rect2[0]+rect2[2]])
    if x1 > x0 and y1 > y0:
        ratio = (x1 - x0) * (y1 - y0) / (min(rect1[2], rect2[2]) * min(rect1[3], rect2[3]))
    else:
        ratio = -1
    return ratio, [x0, y0, x1 - x0, y1 - y0]
    


def debug_draw_img(bgr_img, json_data):
    draw_img = copy.copy(bgr_img)
    colors = [
        (0,0,128),
        (0,0,255),
        (0,128,0),
        (0,128,128),
        (0,128,255),
        (0,255,0),
        (0,255,128),
        (0,255,255),

        (128,0,0),
        (128,0,128),
        (128,0,255),
        (128,128,0),
        (128,128,128),
        (128,128,255),
        (128,255,0),
        (128,255,128),
        (128,255,255),

        (255,0,0),
        (255,0,128),
        (255,0,255),
        (255,128,0),
        (255,128,128),
        (255,128,255),
        (255,255,0),
        (255,255,128),

    ]
    
    save_dir = "./api_yolov8_det_img"
    if not osp.exists(save_dir):
        os.makedirs(save_dir, exist_ok=True)
    
    
    # remove_seal_img = remove_red_hsv(bgr_img)

    cur_time = time.time()
    for idx, item in enumerate(json_data["data"]):
        x, y, w, h = item["position"]
        category = item["category"]
        confidence = item["confidence"]
        class_idx = item["class_idx"]
        
        # # unclip
        # w += min(20, x)
        # x -= min(20, x)
        # h += min(20, y)
        # y -= min(20, y)
        # w += min(20, draw_img.shape[1]-x-w)
        # h += min(20, draw_img.shape[0]-y-h)

        # 保存原始图片
        if category in ["date"]:
            if h < 28:
                h += min(y, 2)
                y -= min(y, 2)
                if y + h > bgr_img.shape[0]:
                    h = bgr_img.shape[0] - y
        crop_img = bgr_img[y:y+h, x:x+w]
        cv2.imwrite(osp.join(save_dir, "{}.crop.{}.{}.jpg".format(cur_time, category, idx)), crop_img)
        

        # 保存去红之后的图片
        if category in ["handwrite", "date", "pagenum"]:
            crop_img_rmred = remove_red_handwrite_date_hsv_2(crop_img)
            cv2.imwrite(osp.join(save_dir, "{}.crop.{}.{}.rmred.jpg".format(cur_time, category, idx)), crop_img_rmred)
        
        if category in ["stamp"]:
            extract_stamp, bin_img = remove_black_seal_hsv(crop_img)
            # extract_stamp, bin_img = extract_red_seal_hsv(crop_img)
            cv2.imwrite(osp.join(save_dir, "{}.crop.{}.{}.extract.jpg".format(cur_time, category, idx)), extract_stamp)
        
        cv2.rectangle(draw_img, (x,y), (x+w, y+h), colors[class_idx%len(colors)], 2)
        cv2.putText(draw_img,
                    category + " " + str(confidence)[:4],
                    (x,y-5),
                    cv2.FONT_HERSHEY_COMPLEX,
                    1,
                    colors[class_idx%len(colors)],
                    2)
    
    cv2.imwrite(osp.join(save_dir, "{}.draw.jpg".format(cur_time)), draw_img)


def api_yolov8_det_img(bgr_img, img_b64, filter_score=0, url="", debug=0):
    headers = {'Content-Type': 'application/json;charset=UTF-8'}
    
    data = 	{
        "img_base64" : img_b64
    }
    
    headers = {'Content-Type': 'application/json;charset=UTF-8'}
    string = json.dumps(data, cls=MyEncoder)
    res = requests.post(url, data=string, headers=headers)
    
    if res.status_code != 200:
        return None
    
    json_data = res.json()
    if filter_score:
        # 内置过滤条件
        new_data = []
        for item in json_data["data"]:
            category = item["category"]
            confidence = item["confidence"]
            if category == "stamp" and confidence < 0.70:
                continue
            if category == "handwrite" and confidence < 0.70:
                continue
            if category == "date" and confidence < 0.7:
                continue
            if category == "pagenum" and confidence < 0.7:
                continue
            new_data.append(item)
        json_data["data"] = new_data
        
    if debug:
        debug_draw_img(bgr_img, json_data)
        
    return json_data

def extract_red_seal_hsv(bgr_img):
    # 章印提取红色部分
    bgr_img_copy = copy.copy(bgr_img)
    hsv_img = cv2.cvtColor(bgr_img_copy, cv2.COLOR_BGR2HSV)
    mask1 = cv2.inRange(hsv_img, np.array((156, 43, 43)), np.array((180, 255, 255)))
    mask2 = cv2.inRange(hsv_img, np.array((0, 43, 43)), np.array((10, 255, 255)))
    mask1 += mask2
    
    background_img = np.zeros(bgr_img_copy.shape, np.uint8) + np.array([255,255,255], np.uint8)
    # new_mat = cv2.cvtColor(mask1, cv2.COLOR_GRAY2BGR)
    # bgr_img_copy = cv2.add(bgr_img_copy, new_mat)
    cv2.copyTo(bgr_img_copy, mask1, background_img)
    
    return background_img, mask1

def remove_black_seal_hsv(bgr_img):
    # 章印去除黑色部分
    bgr_img_copy = copy.copy(bgr_img)
    
    gray_img = cv2.cvtColor(bgr_img_copy, cv2.COLOR_BGR2GRAY)
    bin_img = cv2.threshold(gray_img, 40, 255, cv2.THRESH_BINARY_INV)[1]
    
    # hsv_img = cv2.cvtColor(bgr_img_copy, cv2.COLOR_BGR2HSV)
    # mask1 = cv2.inRange(hsv_img, np.array((156, 43, 43)), np.array((180, 255, 255)))
    # mask2 = cv2.inRange(hsv_img, np.array((0, 43, 43)), np.array((10, 255, 255)))
    # mask1 += mask2
    
    background_img = np.zeros(bgr_img_copy.shape, np.uint8) + np.array([255,255,255], np.uint8)
    # new_mat = cv2.cvtColor(mask1, cv2.COLOR_GRAY2BGR)
    # bgr_img_copy = cv2.add(bgr_img_copy, new_mat)
    cv2.copyTo(background_img, bin_img, bgr_img_copy)
    
    return bgr_img_copy, bin_img


def remove_red_hsv(bgr_img):
    # bgr_img 整张图
    bgr_img_copy = copy.copy(bgr_img)
    hsv_img = cv2.cvtColor(bgr_img_copy, cv2.COLOR_BGR2HSV)
    mask1 = cv2.inRange(hsv_img, np.array((156, 86, 86)), np.array((180, 255, 255)))
    mask2 = cv2.inRange(hsv_img, np.array((0, 86, 86)), np.array((10, 255, 255)))
    mask1 += mask2
    new_mat = cv2.cvtColor(mask1, cv2.COLOR_GRAY2BGR)
    bgr_img_copy = cv2.add(bgr_img_copy, new_mat)
    return bgr_img_copy

def remove_red_handwrite_date_hsv_2(small_img):
    # handwrite date 小图专用
    bgr_img_copy = copy.copy(small_img)
    
    gray_img = cv2.cvtColor(bgr_img_copy, cv2.COLOR_BGR2GRAY)
    bin_img = cv2.threshold(gray_img, 100, 255, cv2.THRESH_BINARY_INV)[1]
    # cv2.imwrite("bin.jpg", bin_img)
    
    hsv_img = cv2.cvtColor(bgr_img_copy, cv2.COLOR_BGR2HSV)
    mask1 = cv2.inRange(hsv_img, np.array((156, 43, 43)), np.array((180, 255, 255)))
    mask2 = cv2.inRange(hsv_img, np.array((0, 43, 43)), np.array((10, 255, 255)))
    mask1 += mask2
    mask1 = cv2.subtract(mask1, bin_img)
    
    new_mat = cv2.cvtColor(mask1, cv2.COLOR_GRAY2BGR)
    bgr_img_copy = cv2.add(bgr_img_copy, new_mat)
    return bgr_img_copy


def pack_yolov8_handwrite_data(categories, bgr_img, ret_data_handwrite, debug_mode):
    for idx, item in enumerate(ret_data_handwrite["data"]):
        x, y, w, h = item["position"]
        category = item["category"]
        confidence = item["confidence"]
        
        # 整体去红
        # remove_seal_img = remove_red_hsv(bgr_img)
        
        if category == "stamp":
            crop_img = bgr_img[y:y+h, x:x+w]
            # TODO
            # 目前测试两种思路效果都不好
            # # seal_img, bin_img = extract_red_seal_hsv(crop_img)
            # seal_img, bin_img = remove_black_seal_hsv(crop_img)
            # seal_data = seal_recognition_func(seal_img, debug_mode)
            # 选择用原图识别
            seal_data = seal_recognition_func(crop_img, debug_mode)
            seal_data["position"] = [[x,y], [x+w,y], [x+w,y+h], [x,y+h]]
            seal_data["det_score"] = confidence
            
            seal_multicls_data = seal_class_func(crop_img, debug_mode)
            seal_data["cls_data"] = seal_multicls_data
            
            # 章印重叠默认为0
            seal_data["seal_cover"] = 0
            # 计算iou
            for index, tmp_data in enumerate(ret_data_handwrite["data"]):
                if tmp_data["category"] != "stamp": 
                    continue
                if index == idx:
                    continue
                if iou_3([x, y, w, h], tmp_data["position"])[0] > 0.1:
                    seal_data["seal_cover"] = 1
                    
            categories["seals"].append(seal_data)
            
        if category == "handwrite":
            crop_img = bgr_img[y:y+h, x:x+w]
            crop_img_rmred = remove_red_handwrite_date_hsv_2(crop_img)
            hw_data = handwrite_rec_func(crop_img_rmred, debug_mode)
            hw_data["position"] = [[x,y], [x+w,y], [x+w,y+h], [x,y+h]]
            hw_data["det_score"] = confidence
            categories["handwritten_signatures"].append(hw_data)
            
        if category == "date":
            if h < 28:
                h += min(y, 2)
                y -= min(y, 2)
                if y + h > bgr_img.shape[0]:
                    h = bgr_img.shape[0] - y
            crop_img = bgr_img[y:y+h, x:x+w]
            crop_img_rmred = remove_red_handwrite_date_hsv_2(crop_img)
            date_data = handwrite_rec_func(crop_img_rmred, debug_mode)
            date_data["position"] = [[x,y], [x+w,y], [x+w,y+h], [x,y+h]]
            date_data["det_score"] = confidence
            categories["signature_dates"].append(date_data)
            

def pack_yolov8_pagenum_data(categories, bgr_img, ret_data_date, debug_mode):
    for item in ret_data_date["data"]:
        x, y, w, h = item["position"]
        crop_img = bgr_img[y:y+h, x:x+w]
        category = item["category"]
        confidence = item["confidence"]
        
        if category == "pagenum":
            pagenum_data = pagenum_rec_func(crop_img, debug_mode)
            pagenum_data["position"] = [[x,y], [x+w,y], [x+w,y+h], [x,y+h]]
            pagenum_data["det_score"] = confidence
            categories["page_number"] = pagenum_data
            break
            
        
if __name__ == "__main__":

    seal_url = "http://183.221.0.158:29986/api/algo/docLayout"
    pagenum_url = "http://183.221.0.158:29986/api/algo/pageNum"

    file_path = sys.argv[1]
    file_paths = []
    if osp.isfile(file_path):
        file_paths.append(file_path)
    else:
        file_paths = filesystem.get_all_filepath(file_path, [".jpg", ".png"])
    
    for file_path in file_paths:
        bgr_img = cv2.imread(file_path)
        buffer = io.BytesIO()
        bgr_img = cv2.imread(file_path)
        rgb_img = bgr_img[:, : , ::-1]
        pil_image = Image.fromarray(rgb_img) 
        pil_image.save(buffer, format="jpeg")
        img_b64 = base64.b64encode(buffer.getvalue())

        debug = 1
        filter_score = 1
        yolov8_ret_data = api_yolov8_det_img(bgr_img, img_b64, filter_score=filter_score, url=seal_url, debug=debug)
        # yolov8_ret_data = api_yolov8_det_img(bgr_img, img_b64, filter_score=filter_score, url=pagenum_url, debug=debug)

        # 用于测试章印去除黑色效果
        # 目前测试两种思路效果都不好
        # # seal_img, bin_img = extract_red_seal_hsv(bgr_img)
        # seal_img, bin_img = remove_black_seal_hsv(bgr_img)
        # cv2.imwrite(file_path+".bin.jpg", bin_img)
        # cv2.imwrite(file_path+".seal.jpg", seal_img)
